from keras.applications.resnet import ResNet50
from tensorflow.keras.layers import (Input)
from tensorflow.keras.models import Model

from model import conv_layer


def build_backbone(input_shape,
                   n_layers=4, weights='imagenet'):
    inputs = Input(shape=input_shape)

    features = ResNet50(include_top=False, weights=weights)(inputs)

    outputs = [features]

    n_filters = 64
    prev_conv = features
    for i in range(n_layers - 1):
        postfix = "_layer" + str(i + 2)
        conv = conv_layer(prev_conv,
                          n_filters,
                          kernel_size=3,
                          strides=2,
                          use_maxpool=False,
                          postfix=postfix)
        outputs.append(conv)
        prev_conv = conv
        n_filters *= 2

    # instantiate model
    model = Model(inputs=inputs,
                  outputs=outputs,
                  name='ssd_resnet50_pretrained')
    return model
